import uuid
import io
from flask import Flask, render_template, request, jsonify
from typing import Optional, List
from pathlib import Path
from flask_sqlalchemy import SQLAlchemy
from minio import Minio
#from chromadb.config import Settings
import chromadb
from sentence_transformers import SentenceTransformer
from text_splitter import TextSplitter, SplitConfig

# ------------------------------
# Flask 与 SQLAlchemy 初始化
# ------------------------------
app = Flask(__name__)
# 这里使用 MySQL 数据库，注意替换 user、password、host、port 和 db 名称
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://root:123456@192.168.202.128:3306/demo'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)

# 定义文档和文档片段元数据的数据模型
class Document(db.Model):
    __tablename__ = 'documents'
    id = db.Column(db.String(36), primary_key=True)  # 使用 UUID
    filename = db.Column(db.String(255), nullable=False)
    minio_path = db.Column(db.String(255), nullable=False)
    # 你可以根据需要增加其它字段，如上传时间、文件类型等

class DocumentSegment(db.Model):
    __tablename__ = 'document_segments'
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    document_id = db.Column(db.String(36), db.ForeignKey('documents.id'), nullable=False)
    segment_index = db.Column(db.Integer, nullable=False)
    text = db.Column(db.Text, nullable=False)
    chroma_collection = db.Column(db.String(255), nullable=False)
    chroma_id = db.Column(db.String(255), nullable=False)
    # 其他字段也可以添加

# ------------------------------
# MinIO 客户端初始化
# ------------------------------
# 替换下方的地址和访问密钥
minio_client = Minio(
    "192.168.202.128:9000",
    access_key="MOV5r2pvqHG1Irjj6kPj",
    secret_key="3jE2hPSciG6xeiYP5qWIG0b8kVPY4b5eMVoSCTO8",
    secure=False
)
bucket_name = "documents"
if not minio_client.bucket_exists(bucket_name):
    minio_client.make_bucket(bucket_name)

# ------------------------------
# ChromaDB 客户端初始化
# ------------------------------
chroma_client = chromadb.PersistentClient(path=str(Path("./chroma_db").resolve()))

collection_name = "document_segments"
vector_collection = chroma_client.get_or_create_collection(name=collection_name)

# ------------------------------
# 嵌入模型初始化
# ------------------------------
# 使用 SentenceTransformer 生成文本嵌入
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

@app.route('/upload_page')
def upload_page():
    return render_template('upload.html')

# ------------------------------
# API 接口：上传文档并处理
# ------------------------------
@app.route('/upload_document', methods=['POST'])
def upload_document():
    # 检查文件是否存在
    if 'file' not in request.files:
        return jsonify({"error": "未提供文件"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "文件名为空"}), 400

    # 获取分割策略参数
    strategy = request.form.get('strategy', 'paragraph')  # 获取切分策略
    chunk_size = int(request.form.get('chunk_size', 1000))
    chunk_overlap = int(request.form.get('chunk_overlap', 200))

    # 为文档生成唯一 ID
    document_id = str(uuid.uuid4())
    filename = file.filename
    object_name = f"{document_id}/{filename}"

    # 读取文件内容
    data = file.read()
    if not data:
        return jsonify({"error": "文件为空"}), 400

    # 将文件上传到 MinIO
    file_data = io.BytesIO(data)
    file_size = len(data)
    try:
        minio_client.put_object(bucket_name, object_name, file_data, file_size)
    except Exception as e:
        return jsonify({"error": "MinIO 存储失败", "details": str(e)}), 500

    # 将文档元数据保存到 MySQL
    doc = Document(id=document_id, filename=filename, minio_path=object_name)
    db.session.add(doc)
    db.session.commit()

    # 将文档内容解析为段落
    try:
        content = data.decode('utf-8')
    except Exception as e:
        return jsonify({"error": "文件解码失败", "details": str(e)}), 400

    # 配置切分策略
    split_config = SplitConfig(
        strategy=strategy,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap
    )

    segments = process_document(content, split_config)

    if not segments:
        segments = [content]  # 若分段失败，则作为一个整体

    # 计算每个段落的嵌入向量
    embeddings = embedding_model.encode(segments).tolist()

    # 逐段保存到向量数据库并存储段落元数据
    for index, (para, emb) in enumerate(zip(segments, embeddings)):
        # 为段落生成一个唯一的 ID（可以根据需要自定义）
        segment_id = f"{document_id}_{index}"
        try:
            # 将嵌入和对应的文本、元数据添加到 ChromaDB
            vector_collection.add(
                ids=[segment_id],
                embeddings=[emb],
                metadatas=[{"document_id": document_id, "segment_index": index}],
                documents=[para]
            )
        except Exception as e:
            return jsonify({"error": "向量数据库存储失败", "details": str(e)}), 500

        # 保存文档片段元数据到 MySQL
        segment = DocumentSegment(
            document_id=document_id,
            segment_index=index,
            text=para,
            chroma_collection=collection_name,
            chroma_id=segment_id
        )
        db.session.add(segment)
    db.session.commit()

    return jsonify({
        "message": "文档上传及处理成功",
        "document_id": document_id,
        "num_segments": len(segments)
    })
def process_document(content: str, split_config: Optional[SplitConfig] = None) -> List[str]:
    if split_config is None:
        split_config = SplitConfig()  # 使用默认配置
    
    splitter = TextSplitter(split_config)
    return splitter.split(content)

# ------------------------------
# 其它 API 接口（可按需扩展）
# ------------------------------
# 例如：单独保存文档元数据、单独保存文档片段等
# 此处仅给出一个完整的上传示例

if __name__ == '__main__':
    # 确保数据库表已经创建（仅用于测试，生产环境建议使用数据库迁移工具如 Alembic）
    with app.app_context():
        db.create_all()
    app.run(debug=True)
